import numpy as np
from sklearn.metrics import precision_recall_fscore_support
from trade.model_base import FitBase


class RollingFit(FitBase):
    def __init__(self, model):
        self.model = model  # classify

    def fit(self, X, y, rolls: int = 10, other: int = 1):
        """
        :param X:
        :param y:
        :param rolls:
        :param other: 默认的不分类
        :return:
        """
        old_y = y
        y_cls = np.copy(y)
        # mask = np.random.choice([False, True], size=n, p=[2 / 3, 1 / 3])
        for i in range(rolls):
            unique_values, counts = np.unique(y_cls, return_counts=True)
            weights = {k: len(y_cls) / (v + 1) for k, v in zip(unique_values, counts)}
            weight_array = np.zeros(max(weights.keys()) + 1)
            for k, v in weights.items():
                weight_array[k] = v
            sample_weight = weight_array[y_cls]
            print(
                f"start fit cls counts {len(y_cls)}: {np.unique(y_cls, return_counts=True)}"
            )
            print(f"Fit roll cur weights: {weights}")
            self.model.fit(X, y_cls, sample_weight=sample_weight)
            results = self.model.predict(X)
            results = results.squeeze(-1)
            precision, recall, f1, support = precision_recall_fscore_support(
                old_y[results != other], results[results != other], average=None
            )
            print(
                f"Fit roll pred model {i} precision recall f1 support {precision} {recall} {f1} {support}"
            )
            y_cls = np.where(results == old_y, results, other)

    def predict(self, X):
        results = self.model.predict(X)
        return results

    def save_model(self, file: str):
        self.model.save_model(file)

    def load_model(self, file: str):
        self.model.load_model(file)
